from utils import *
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import argparse
import os
from scipy.optimize import linear_sum_assignment as linear_assignment
from PCA import PCA
from K_means import K_Means

transform = transforms.Compose([
    transforms.ToTensor(),
])

dataset = None
batch_size = None
dataloader = None
device = None


def main(args):
    clust = None
    if args.eval:
        for images, labels in dataloader:
            images = images.reshape(-1, 28 * 28).squeeze(1).to(args.device)
            print('Applying K-Means...')
            clust = K_Means(args.num_clust, args.device)
            clust.cluster(images, 100)
            if args.num_clust == 10:
                print('Accuracy:', clust.ACC(labels))
            print('----------------------------|')
            print('Applying PCA...')
            print('Size of dataset:', images.size())
            pca = PCA(args.n_components, 'cpu')
            pca.fit(images)
            afterPCA = pca.transform(images)
            print('Size after PCA:', afterPCA.size())
            print('----------------------------|')
            print('Applying K-Means...')
            clust = K_Means(args.num_clust, args.device)
            clust.cluster(images, 100)
            if args.num_clust == 10:
                print('Accuracy:', clust.ACC(labels))
    
    if args.draw:
        print('----------------------------|')
        for images, labels in dataloader:
            images = images.reshape(-1, 28 * 28).squeeze(1).to(args.device)
            print('Applying K-Means...')
            clust = K_Means(args.num_clust, args.device)
            clust.cluster(images, 100)

            print('----------------------------|')
            print('Applying PCA to get 2d results of clustering...')

            pca = PCA(2, 'cpu')
            pca.fit(images)
            afterPCA = pca.transform(images)

            show_clusters(f'results/clusters_K_{args.num_clust}.png', afterPCA[:1000], clust.ind[:1000], args.num_clust)
        


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Program for pr3!!!')
    parser.add_argument('--draw', action='store_true', help='if we show images in one picture')
    parser.add_argument('--eval', action='store_true', help='if we evaluate ACCURACY with dimension before and after PCA')
    parser.add_argument('--dataset', default='../../../Diffusion_Classifier_MNIST/data', type=str, help='ur path 4 MNIST')
    parser.add_argument('--device', default='cpu', type=str, help='device for pytorch')
    parser.add_argument('--num_clust', default=10, type=int, help='specified count of clusters')
    parser.add_argument('--n_components', default=100, type=int, help='n_component for PCA during evaluation')
    args = parser.parse_args()

    assert(os.path.exists(args.dataset))
    print('Dataset loading...')
    dataset = datasets.MNIST(root= args.dataset, train=False, download=False, transform=transform)
    batch_size = len(dataset)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
    device = args.device
    print('Loading finished!!!')

    main(args)